#!/usr/bin/python
import sys, re
import tm2Semmler 

def read_file(corpus, split):
    """ Read and split file into paragraphs
    """
    try:
        f = open(corpus, 'r')
        text = f.read()
        text = re.split(split, text)
        text = text[:-1] # improve!
        f.close()
    except IOError:
        print "IO ERROR:", corpus, "could not be read."
        sys.exit(1)
    return text

def add_signs(text, n): # improve!
    """ add start and stop sign to every paragraph
    """
    text2 = []

    for paragraph in text:
        text2.append(['START']*(n-1) + paragraph.split() + ['STOP'])
    
    return text2

def merge_dict(bigDict, smallDict):
    """ merges two dicts. In case of any collsiding keys, the value of both is summed
    """
    for k in smallDict:
        if k in bigDict:
            bigDict[k] = smallDict[k] + bigDict[k]
        else:
            bigDict[k] = smallDict[k]
    return bigDict

def n_gram(text, n):
    """  create n-grams from word list
    """
    nGrams = {} 
    for paragraph in text:
        nGram = tm2Semmler.n_gram(paragraph, n)
        nGrams = merge_dict(nGrams, nGram)
    return nGrams

def filter_sen(sentences, n):
    """ filter all sentences whose length is unequal to n
    """
    newSentences = []

    for sentence in sentences:
        if len(sentence.split(' ')) == n:
            newSentences += [sentence]
    return newSentences

def sub_word(sen):
    """ subtracts the last word of a sentence (string)
    """
    return ' '.join(sen.split(' ')[:-1])

def calc_prob_sen(sentences, nGrams, nMinGrams):
    """ Count(w1, ..., wi) * Count(w1, ..., wi-1)
    """
    results = []
    for sentence in sentences:
        if sentence in nGrams:
            results += [float(nGrams[sentence]) / nMinGrams[sub_word(sentence)]]
        else:
            results += [-1]
    return results

def calc_prob_sen2(sentences, nGrams, nMinGrams, n):
    """ P(wi | wi-n+1, ..., wi-1)
    """
    results = []
    for sentence in sentences:
        sen = sentence.split(' ')
        leng = len(sen)
        newSentences = []

        for i in range(1,leng+1):
            if i-n<0:
                newSentences += ['START ' * abs(i-n) + ' '.join(sen[:i])]
            else:
                newSentences += [' '.join(sen[i-n:i])]
        newSentences += [' '.join(sen[-n+1:]) + ' STOP']

        results += [reduce(lambda x,y: x*y, calc_prob_sen(newSentences, nGrams, nMinGrams))]
    return results

def perm(a):
    """ permutation
    """
    if len(a) <= 1:
        yield a
    else:
        for el in range(0, len(a)):
            for rest in perm(a[:el] + a[el+1:]):
                yield [a[el]] + rest 

def calc_prob_of_perm(someSet, nGrams, nMinGrams, n):
    """ the highest prob of all permuations
    """
    prob = {}
    for onePerm in perm(someSet):
        sen = ' '.join(onePerm)
        prob[sen] = float(calc_prob_sen2([sen], nGrams, nMinGrams, n)[0])
    maxKey = max(prob, key=prob.get)
    return maxKey, prob[maxKey]

def question_1(text):
    print "Question 1"
    text2 = add_signs(text, 2)
    print tm2Semmler.freq_list(n_gram(text2,2))[0:10]

def question_2(addText, nGrams, nMinGrams):
    print "Question 2"
    print addText
    print calc_prob_sen(addText, nGrams, nMinGrams)

def question_3(addText2, nGrams, nMinGrams, n):
    print "Question 3"
    print addText2
    print calc_prob_sen2(addText2, nGrams, nMinGrams, n)

def question_4(nGrams, nMinGrams, n):
    print "Question 4"
    setA = ['She', 'daughters', 'youngest', 'was', 'of', 'the', 'two']
    setB = ['She', 'was', 'the', 'youngest']
    print 'set A: ', setA
    print calc_prob_of_perm(setA, nGrams, nMinGrams, n)
    print 'set B: ', setB
    print calc_prob_of_perm(setB, nGrams, nMinGrams, n)

def main(argv=sys.argv):
    corpus  = argv[1]
    n       = int(argv[2])
    addFile = argv[3]
    addFile2 = argv[4]

    text = read_file(corpus, '\n[\n]+')
    nText = add_signs(text, n)
    nMinText = add_signs(text, n)

    addText = read_file(addFile, '\n')
    addText = filter_sen(addText, n)
    
    addText2 = read_file(addFile2, '\n')
    
    nGrams    = n_gram(nText, n)
    nMinGrams = n_gram(nMinText, n-1)
    nMinFreqList = tm2Semmler.freq_list(nMinGrams)

    question_1(text)
    question_2(addText, nGrams, nMinGrams)
    question_3(addText2, nGrams, nMinGrams, n)
    question_4(nGrams, nMinGrams, n)
    return 0

if __name__ == "__main__":
    sys.exit(main())
